import argparse
import collections
import datetime
import glob
import json
import os.path as osp
import sys
import uuid
import time

import imgviz
import numpy as np

import labelme

try:
    import pycocotools.mask
except ImportError:
    print("Please install pycocotools:\n\n    pip install pycocotools\n")
    sys.exit(1)


# https://github.com/pascal1129/kaggle_airbus_ship_detection/blob/master/0_rle_to_coco/1_ships_to_coco.py

def main():

    # 首先生成label文件
    with open('labels.txt', 'w') as f:  # 设置文件对象
        f.write("__ignore__")
        f.write("\n")
        f.write("__background__")
        f.write("\n")
        f.write("cron")  # 将字符串写入文件中
        f.write("\n")

    parser = argparse.ArgumentParser(description="json2coco")
    # 原始json文件保存的路径
    parser.add_argument("--input_dir", help="input annotated directory",
                        default="")
    # 整合后的json文件保存的路径
    parser.add_argument("--output_dir", help="output dataset directory",
                        default="coco2017/annotations")
    parser.add_argument("--labels", help="labels file",
                        default='labels.txt')  # required=True
    args = parser.parse_args()

    now = datetime.datetime.now()
    start = time.time()

    data = dict(
        info=dict(
            description="cron datasets",
            url=None,
            version="label=4.5.6",
            year=now.year,
            contributor="David",
            date_created=now.strftime("%Y-%m-%d %H:%M:%S.%f"),
        ),
        # licenses=[dict(url=None, id=0, name=None,)],
        images=[
            # license, url, file_name, height, width, date_captured, id
        ],
        type="instances",
        annotations=[
            # segmentation, area, iscrowd, image_id, bbox, category_id, id
        ],
        categories=[
            # supercategory, id, name
        ],
    )

    class_name_to_id = {}
    for i, line in enumerate(open(args.labels).readlines()):
        class_id = i - 1  # starts with -1
        class_name = line.strip()
        if class_id == -1:
            assert class_name == "__ignore__"
            continue
        if class_id == 0:
            assert class_name == "__background__"
            continue
        class_name_to_id[class_name] = class_id
        # print(class_id,class_name,'\n')
        data["categories"].append(
            dict(supercategory="cron", id=class_id, name=class_name, )  # 一类目标+背景，id=0表示背景
        )
    print("categories 生成完成", '\n')

    out_ann_file = osp.join(args.output_dir, "instances_train2017.json")  # 自动添加"/

    label_files = glob.glob(osp.join(args.input_dir, "*.json"))  # 图像id从json文件中读取
    for image_id, filename in enumerate(label_files):
        print(image_id, filename)
        # print("Generating dataset from:", filename)

        label_file = labelme.LabelFile(filename=filename)

        base = osp.splitext(osp.basename(filename))[0]  # 图片名
        out_img_file = osp.join(args.output_dir, base + ".jpg")  # 保存图片路径

        img = labelme.utils.img_data_to_arr(label_file.imageData)
        imgviz.io.imsave(out_img_file, img)
        data["images"].append(
            dict(
                # license=0,
                # url=None,
                file_name=osp.relpath(out_img_file, osp.dirname(out_ann_file)),
                height=img.shape[0],
                width=img.shape[1],
                # date_captured=None,
                id=image_id,
            )
        )

        masks = {}  # for area
        segmentations = collections.defaultdict(list)  # for segmentation
        for shape in label_file.shapes:
            points = shape["points"]
            label = shape["label"]
            group_id = shape.get("group_id")
            shape_type = shape.get("shape_type", "polygon")
            mask = labelme.utils.shape.shape_to_mask(img.shape[:2], points, shape_type)  # labelme=4.5.6的shape_to_mask函数
            if group_id is None:
                group_id = uuid.uuid1()

            instance = (label, group_id)
            # print(instance)

            if instance in masks:
                masks[instance] = masks[instance] | mask
            else:
                masks[instance] = mask

            if shape_type == "rectangle":
                (x1, y1), (x2, y2) = points
                x1, x2 = sorted([x1, x2])
                y1, y2 = sorted([y1, y2])
                points = [x1, y1, x2, y1, x2, y2, x1, y2]
            else:
                points = np.asarray(points).flatten().tolist()

            segmentations[instance].append(points)
        segmentations = dict(segmentations)

        for instance, mask in masks.items():
            cls_name, group_id = instance
            #            if cls_name not in class_name_to_id:
            #                continue
            #            cls_id = class_name_to_id[cls_name]

            mask = np.asfortranarray(mask.astype(np.uint8))

            mask = pycocotools.mask.encode(mask)

            area = float(pycocotools.mask.area(mask))
            bbox = pycocotools.mask.toBbox(mask).flatten().tolist()

            data["annotations"].append(
                dict(
                    id=len(data["annotations"]),
                    image_id=image_id,
                    category_id=1,  # 都是1类cls_id
                    segmentation=segmentations[instance],
                    area=area,
                    bbox=bbox,
                    iscrowd=0,
                )
            )

    print("annotations 生成完成", '\n')

    with open(out_ann_file, "w") as f:
        json.dump(data, f, indent=2)

    cost_time = (time.time() - start) / 1000
    print("cost_time:{:.2f}s".format(cost_time))


if __name__ == "__main__":
    main()

